# -*- coding: utf-8 -*-


import mitmproxy.http
from mitmproxy import ctx
import time
import datetime
# import xlwt
import sqlite3

# 保存到数据库
conn = sqlite3.connect('http_connet.db')
# 只记录特定特征的网址
# url_key_word = "zljllcy.com"
# url_key_word = "*"
url_key_word = "hkitcloud.net"

class Saver:
    def __init__(self):
        self.num = 0
        self.requestNum = 0
        self.responseOrErrorNum = 0
        self.aa = 0
        # 特定请求域名才记录
        self.url_key_word = url_key_word
        # 是否记录
        self.mark_flag = 1
        self.all_arr = [[]]
        # 链接记录的字典
        self.connetc_dict = {}
        self.flow_tmp = []
        self.request_time = ""
        # 请求记录
        c = conn.cursor()
        c.execute('''CREATE TABLE IF NOT EXISTS request_info
               (ID  INTEGER PRIMARY KEY  NOT NULL,
               flow_id       CHAR(50)          ,
               URL           TEXT      NOT NULL,
               HOST          TEXT      NOT NULL,
               Path          TEXT      NOT NULL,
               RequestType   TEXT      ,
               RequestData   TEXT      ,
               RequestTime   CHAR(25) ,
               RecodTime     CHAR(25));''')

        c.execute('''CREATE TABLE IF NOT EXISTS response_info
               (ID  INTEGER PRIMARY KEY  NOT NULL,
               flow_id       CHAR(50),
               ResponseData  TEXT    ,
               RecodTime     CHAR(25));''')

        c.execute('''CREATE TABLE IF NOT EXISTS httrequest
               (ID  INTEGER PRIMARY KEY  NOT NULL,
               URL           TEXT      NOT NULL,
               HOST          TEXT      NOT NULL,
               Path          TEXT      NOT NULL,
               RequestType   TEXT      ,
               RequestData   TEXT      ,
               ResponseData  TEXT      ,
               RecodTime     CHAR(25)  ,
               RequestTime   CHAR(25));''')# '请求路径', '请求域名', '请求path', '请求类型', '请求报文' ,'响应报文', '记录时间'

        conn.commit()
        c.close()

    def http_connect(self, flow: mitmproxy.http.HTTPFlow):
        """
        收到了来自客户端的 HTTP CONNECT 请求，此时只是客户端到mitmproxy的请求，还没触发到server的请求，所以此时还不会触发request等事件
        可以在此时对请求地址之类的做操作
        """
        pass
        # self.flow_tmp = []
        # tmp = [str(datetime.datetime.now()), '||', 'connect', '||', str(flow.id), '||', flow.type]
        # self.log_out(tmp)

    def request(self, flow: mitmproxy.http.HTTPFlow):
        """
        自客户端的 HTTP 请求被成功【完整】读取。
        """
        # tmp = [str(datetime.datetime.now()), '||', 'request', '||', str(flow.id), '||', flow.request.url]
        # self.log_out(tmp)
        self.mark_flag = 1
        print("in request")
        tmp_dict = {}
        # 请求的不是要抓取的地址
        if self.url_key_word == "*":
            self.mark_flag = 1
        elif self.url_key_word not in flow.request.host:
            self.mark_flag = 0
        # 请求的是 js脚本、图片、音频等信息
        rq_type = flow.request.path.split('.')[-1]
        if rq_type in ['js', 'png', 'jpeg', 'jpg', 'gif', 'html', 'apk']:
            self.mark_flag = 0
        # 请求的类型是option的不记录
        if flow.request.method == 'OPTIONS':
            self.mark_flag = 0
        # 如果不过滤地址，则完整记录所有信息
        if url_key_word == "*":
            self.mark_flag = 1
        # unix时间戳转为本地时间

        if self.mark_flag == 0:
            # 删除字典中的记录
            # self.connetc_dict.pop(str(flow.id))
            pass
            print("opus, i am here, but not should be")
        else:
            print("i am here")
            self.request_time = str(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(flow.request.timestamp_start)))
            # 加入字典
            tmp_dict['request_time'] = self.request_time
            tmp_dict['url'] = flow.request.url
            tmp_dict['host'] = flow.request.host
            tmp_dict['path'] = flow.request.path
            tmp_dict['method'] = flow.request.method
            # 请求报文
            tmp_dict['content'] = str(flow.request.get_content(), 'utf-8')

            cor = conn.cursor()
            cor.execute("INSERT INTO request_info (ID,flow_id,URL,HOST,Path,RequestType,RequestData,RecodTime,RequestTime) \
                                        VALUES (?,?,?,?,?,?,?,?,?) ", \
                        (None, str(flow.id), tmp_dict['url'], tmp_dict['host'], tmp_dict['path'], tmp_dict['method'], tmp_dict['content'], \
                         tmp_dict['request_time'], time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))))
            cor.close()
            conn.commit()




    def error(self, flow: mitmproxy.http.HTTPFlow):
        """发生了一个 HTTP 错误。比如无效的服务端响应、连接断开等"""
        print("http_err", flow.request.url)
        # self.aa = self.aa + 1
        # self.responseOrErrorNum = self.responseOrErrorNum + 1
        # self.flow_tmp.append("Error response")

    def response(self, flow: mitmproxy.http.HTTPFlow):
        """自服务端端的 HTTP 响应被成功完整读取"""
        # tmp = [str(datetime.datetime.now()), '||', 'response', '||', str(flow.id), '||', flow.request.url]
        # self.log_out(tmp)
        if self.mark_flag == 1:
            # 请求大小
            try:
                # self.flow_tmp.append(flow.request.headers['Content-Length'])
                pass
            except:
                # self.flow_tmp.append("")
                pass
            # 响应大小
            try:
                # self.flow_tmp.append(flow.response.headers['Content-Length'])
                pass
            except:
                # self.flow_tmp.append("")
                pass
            # 响应类型
            try:
                # self.flow_tmp.append(flow.response.headers['Content-Type'])
                pass
            except:
                # self.flow_tmp.append("")
                pass
            cor = conn.cursor()
            cor.execute("INSERT INTO response_info (ID,flow_id,ResponseData,RecodTime) \
                                        VALUES (?,?,?,?) ", \
                        (None, str(flow.id), flow.response.text, time.strftime("%Y-%m-%d %H:%M:%S",time.localtime(time.time()))))
            cor.close()
            conn.commit()


    def log_out(self, array):
        return 0
        with open("mitmproxy.log", 'a+') as fp:
            try:
                fp.writelines(array)
                fp.write('\r\n')
            except Exception as bug:
                print("write file fail", bug)

    def save(self, dct):
        # 是否保存到exce
        # print(self.connetc_dict)
        cor = conn.cursor()
        # 保存到数据库
        for k in dct.keys():
            try:
                cor.execute("INSERT INTO httrequest (ID,URL,HOST,Path,RequestType,RequestData,ResponseData,RecodTime,RequestTime) \
                            VALUES (?,?,?,?,?,?,?,?,?) ", \
                (None, dct[k]['url'], dct[k]['host'], dct[k]['path'], dct[k]['method'], dct[k]['content'], \
                 dct[k]['response_txt'], dct[k]['recod_time'], dct[k]['request_time']))
            except Exception as bug:
                ctx.log.info("save datadb faile")
                ctx.log.info(bug)
        conn.commit()

    def done(self):
        """
        addon 关闭或被移除，又或者 mitmproxy 本身关闭。由于会先等事件循环终止后再触发该事件，所以这是一个 addon 可以看见的最后一个事件。
        """
        conn.close()

